#include "nn_algo.hpp"

#include "spdlog/spdlog.h"
#include "yololayer.h"

namespace algo {

// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> LoadWeights(const std::string file) {
  SPDLOG_DEBUG("Loading Weights in : {}", file);
  std::map<std::string, Weights> weightMap;

  // Open weights file
  // TODO
  std::ifstream input(file);
  assert(input.is_open() &&
         "Unable to load weight file. please check if the .wts file path is "
         "right!!!!!!");

  // Read number of weight blobs
  int32_t count;
  input >> count;
  // TODO
  assert(count > 0 && "Invalid weight map file.");

  while (count--) {
    Weights wt{DataType::kFLOAT, nullptr, 0};
    uint32_t size;
    std::string name;

    // Read name and type of blob
    input >> name >> std::dec >> size;
    wt.type = DataType::kFLOAT;
    wt.count = size;

    // Load blob
    uint32_t *val = reinterpret_cast<uint32_t *>(malloc(sizeof(val) * size));
    for (uint32_t x = 0, y = size; x < y; ++x) {
      input >> std::hex >> val[x];
    }
    wt.values = val;

    weightMap[name] = wt;
  }
  SPDLOG_WARN("LoadWeights");
  return weightMap;
}

IScaleLayer *addBatchNorm2d(INetworkDefinition *network,
                            std::map<std::string, Weights> &weightMap,
                            ITensor &input, const std::string &lname,
                            float eps) {
  float *gamma = (float *)weightMap[lname + ".weight"].values;
  float *beta = (float *)weightMap[lname + ".bias"].values;
  float *mean = (float *)weightMap[lname + ".running_mean"].values;
  float *var = (float *)weightMap[lname + ".running_var"].values;
  int len = weightMap[lname + ".running_var"].count;

  float *scval = reinterpret_cast<float *>(malloc(sizeof(float) * len));
  for (int i = 0; i < len; i++) {
    scval[i] = gamma[i] / sqrt(var[i] + eps);
  }
  Weights scale{DataType::kFLOAT, scval, len};

  float *shval = reinterpret_cast<float *>(malloc(sizeof(float) * len));
  for (int i = 0; i < len; i++) {
    shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
  }
  Weights shift{DataType::kFLOAT, shval, len};

  float *pval = reinterpret_cast<float *>(malloc(sizeof(float) * len));
  for (int i = 0; i < len; i++) {
    pval[i] = 1.0;
  }
  Weights power{DataType::kFLOAT, pval, len};

  weightMap[lname + ".scale"] = scale;
  weightMap[lname + ".shift"] = shift;
  weightMap[lname + ".power"] = power;
  IScaleLayer *scale_1 =
      network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
  assert(scale_1);
  SPDLOG_WARN("addBatchNorm2d");

  return scale_1;
}

ILayer *convBlock(INetworkDefinition *network,
                  std::map<std::string, Weights> &weightMap, ITensor &input,
                  int outch, int ksize, int s, int g,
                  const std::string &lname) {
  Weights emptywts{DataType::kFLOAT, nullptr, 0};
  int p = ksize / 3;
  IConvolutionLayer *conv1 =
      network->addConvolutionNd(input, outch, DimsHW{ksize, ksize},
                                weightMap[lname + ".conv.weight"], emptywts);
  assert(conv1);
  conv1->setStrideNd(DimsHW{s, s});
  conv1->setPaddingNd(DimsHW{p, p});
  conv1->setNbGroups(g);
  IScaleLayer *bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0),
                                    lname + ".bn", 1e-3);

  // silu = x * sigmoid
  auto sig =
      network->addActivation(*bn1->getOutput(0), ActivationType::kSIGMOID);
  // assert(sig);
  auto ew = network->addElementWise(*bn1->getOutput(0), *sig->getOutput(0),
                                    ElementWiseOperation::kPROD);
  // assert(ew);
  SPDLOG_WARN("convBlock");
  return ew;
}

ILayer *focus(INetworkDefinition *network,
              std::map<std::string, Weights> &weightMap, ITensor &input,
              int inch, int outch, int ksize, const std::string &lname) {
  ISliceLayer *s1 = network->addSlice(
      input, Dims3{0, 0, 0}, Dims3{inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2},
      Dims3{1, 2, 2});
  ISliceLayer *s2 = network->addSlice(
      input, Dims3{0, 1, 0}, Dims3{inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2},
      Dims3{1, 2, 2});
  ISliceLayer *s3 = network->addSlice(
      input, Dims3{0, 0, 1}, Dims3{inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2},
      Dims3{1, 2, 2});
  ISliceLayer *s4 = network->addSlice(
      input, Dims3{0, 1, 1}, Dims3{inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2},
      Dims3{1, 2, 2});
  ITensor *inputTensors[] = {s1->getOutput(0), s2->getOutput(0),
                             s3->getOutput(0), s4->getOutput(0)};
  auto cat = network->addConcatenation(inputTensors, 4);
  auto conv = convBlock(network, weightMap, *cat->getOutput(0), outch, ksize, 1,
                        1, lname + ".conv");
  SPDLOG_WARN("focus");
  return conv;
}

ILayer *bottleneck(INetworkDefinition *network,
                   std::map<std::string, Weights> &weightMap, ITensor &input,
                   int c1, int c2, bool shortcut, int g, float e,
                   const std::string &lname) {
  auto cv1 = convBlock(network, weightMap, input, (int)((float)c2 * e), 1, 1, 1,
                       lname + ".cv1");
  auto cv2 = convBlock(network, weightMap, *cv1->getOutput(0), c2, 3, 1, g,
                       lname + ".cv2");
  if (shortcut && c1 == c2) {
    auto ew = network->addElementWise(input, *cv2->getOutput(0),
                                      ElementWiseOperation::kSUM);
    return ew;
  }
  SPDLOG_WARN("bottleneck end");
  return cv2;
}

ILayer *C3(INetworkDefinition *network,
           std::map<std::string, Weights> &weightMap, ITensor &input, int c1,
           int c2, int n, bool shortcut, int g, float e,
           const std::string &lname) {
  int c_ = (int)((float)c2 * e);
  auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1");
  auto cv2 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv2");
  ITensor *y1 = cv1->getOutput(0);
  for (int i = 0; i < n; i++) {
    auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0,
                        lname + ".m." + std::to_string(i));
    y1 = b->getOutput(0);
  }

  ITensor *inputTensors[] = {y1, cv2->getOutput(0)};
  auto cat = network->addConcatenation(inputTensors, 2);

  auto cv3 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1,
                       lname + ".cv3");
  SPDLOG_WARN("C3 end");
  return cv3;
}

ILayer *SPP(INetworkDefinition *network,
            std::map<std::string, Weights> &weightMap, ITensor &input, int c1,
            int c2, int k1, int k2, int k3, const std::string &lname) {
  int c_ = c1 / 2;
  auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1");

  auto pool1 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX,
                                     DimsHW{k1, k1});
  pool1->setPaddingNd(DimsHW{k1 / 2, k1 / 2});
  pool1->setStrideNd(DimsHW{1, 1});
  auto pool2 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX,
                                     DimsHW{k2, k2});
  pool2->setPaddingNd(DimsHW{k2 / 2, k2 / 2});
  pool2->setStrideNd(DimsHW{1, 1});
  auto pool3 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX,
                                     DimsHW{k3, k3});
  pool3->setPaddingNd(DimsHW{k3 / 2, k3 / 2});
  pool3->setStrideNd(DimsHW{1, 1});

  ITensor *inputTensors[] = {cv1->getOutput(0), pool1->getOutput(0),
                             pool2->getOutput(0), pool3->getOutput(0)};
  auto cat = network->addConcatenation(inputTensors, 4);

  auto cv2 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1,
                       lname + ".cv2");
  SPDLOG_WARN("SPP end");

  return cv2;
}

std::vector<std::vector<float>> GetAnchors(
    std::map<std::string, Weights> &weightMap, const std::string &lname) {
  std::vector<std::vector<float>> anchors;
  Weights wts = weightMap[lname + ".anchor_grid"];
  int anchor_len = Yolo::CHECK_COUNT * 2;
  for (int i = 0; i < wts.count / anchor_len; i++) {
    auto *p = (const float *)wts.values + i * anchor_len;
    std::vector<float> anchor(p, p + anchor_len);
    anchors.push_back(anchor);
  }
  SPDLOG_WARN("GetAnchors");
  return anchors;
}

IPluginV2Layer *AddYoloLayer(INetworkDefinition *network,
                             std::map<std::string, Weights> &weightMap,
                             const std::string &lname,
                             std::vector<IConvolutionLayer *> dets) {
  auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1");
  auto anchors = GetAnchors(weightMap, lname);
  PluginField plugin_fields[2];
  int netinfo[4] = {Yolo::CLASS_NUM, Yolo::INPUT_W, Yolo::INPUT_H,
                    Yolo::MAX_OUTPUT_BBOX_COUNT};
  plugin_fields[0].data = netinfo;
  plugin_fields[0].length = 4;
  plugin_fields[0].name = "netinfo";
  plugin_fields[0].type = PluginFieldType::kFLOAT32;
  int scale = 8;
  std::vector<Yolo::YoloKernel> kernels;
  for (size_t i = 0; i < anchors.size(); i++) {
    Yolo::YoloKernel kernel;
    kernel.width = Yolo::INPUT_W / scale;
    kernel.height = Yolo::INPUT_H / scale;
    memcpy(kernel.anchors, &anchors[i][0], anchors[i].size() * sizeof(float));
    kernels.push_back(kernel);
    scale *= 2;
  }
  plugin_fields[1].data = &kernels[0];
  plugin_fields[1].length = kernels.size();
  plugin_fields[1].name = "kernels";
  plugin_fields[1].type = PluginFieldType::kFLOAT32;
  PluginFieldCollection plugin_data;
  plugin_data.nbFields = 2;
  plugin_data.fields = plugin_fields;
  IPluginV2 *plugin_obj = creator->createPlugin("yololayer", &plugin_data);
  std::vector<ITensor *> input_tensors;
  for (auto det : dets) {
    input_tensors.push_back(det->getOutput(0));
  }
  auto yolo = network->addPluginV2(&input_tensors[0], input_tensors.size(),
                                   *plugin_obj);
  SPDLOG_WARN("AddYoloLayer");
  return yolo;
}

}  // namespace algo
